#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import json
import yaml
import logging
import socket
import threading
from datetime import datetime
from logging.handlers import RotatingFileHandler
import paho.mqtt.client as mqtt

class TCP2MQTTGateway:
    def __init__(self, config_file='config.yaml'):
        # 加载配置
        self.config_file = config_file
        self.config = self._load_config(config_file)
        
        # 初始化日志
        self._init_logging()
        self.logger = logging.getLogger(__name__)
        
        # 初始化MQTT客户端
        self.mqtt_client = None
        self._init_mqtt()
        
        # 初始化TCP连接映射表
        self.tcp_connections = {}  # 用于GUI显示连接状态
        self.topic_to_device = {}
        self.mapping_lock = threading.RLock()
        
        # 心跳配置
        self.heartbeat_enabled = self.config['heartbeat']['enabled']
        self.heartbeat_interval = self.config['heartbeat']['interval']
        self.heartbeat_packet = bytes.fromhex(self.config['heartbeat']['packet'])
        
        # 设备ID在topic中的位置
        self.device_id_position = self.config['mapping']['device_id_position']
        
        # TCP配置
        self.tcp_remote_host = self.config['tcp']['remote_host']
        self.tcp_remote_port = self.config['tcp']['remote_port']
        self.tcp_connect_timeout = self.config['tcp']['connect_timeout']
        self.tcp_recv_buffer_size = self.config['tcp']['recv_buffer_size']
        
        # 停止标志
        self.running = False
        
        # 心跳线程
        self.heartbeat_thread = None
        
        # 启动监控线程
        self.monitor_thread = threading.Thread(target=self._monitor_connections)
        self.monitor_thread.daemon = True
        self.monitor_thread.start()
        
        logging.info("TCP2MQTT Gateway initialized")
    
    def _load_config(self, config_file):
        """加载配置文件"""
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Failed to load config file: {e}")
            sys.exit(1)
    
    def _init_logging(self):
        """初始化日志配置"""
        log_level = getattr(logging, self.config['logging']['level'], logging.INFO)
        log_file = self.config['logging']['file']
        max_bytes = self.config['logging']['max_bytes']
        backup_count = self.config['logging']['backup_count']
        
        # 配置日志格式
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        
        # 设置root logger
        root_logger = logging.getLogger()
        root_logger.setLevel(log_level)
        
        # 移除所有已存在的handler
        for handler in root_logger.handlers[:]:
            root_logger.removeHandler(handler)
        
        # 添加文件handler
        file_handler = RotatingFileHandler(
            log_file, maxBytes=max_bytes, backupCount=backup_count, encoding='utf-8'
        )
        file_handler.setFormatter(formatter)
        root_logger.addHandler(file_handler)
        
        # 添加控制台handler
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        root_logger.addHandler(console_handler)
        
        return logging.getLogger(__name__)
    
    def _init_mqtt(self):
        """初始化MQTT客户端"""
        mqtt_config = self.config['mqtt']
        client_id = mqtt_config['client_id']
        
        # 创建MQTT客户端
        self.mqtt_client = mqtt.Client(client_id=client_id, clean_session=True)
        
        # 设置回调函数
        self.mqtt_client.on_connect = self._on_mqtt_connect
        self.mqtt_client.on_disconnect = self._on_mqtt_disconnect
        self.mqtt_client.on_message = self._on_mqtt_message
        
        # 设置用户名和密码
        if mqtt_config['username'] and mqtt_config['password']:
            self.mqtt_client.username_pw_set(mqtt_config['username'], mqtt_config['password'])
    
    def _on_mqtt_connect(self, client, userdata, flags, rc):
        """MQTT连接成功回调"""
        if rc == 0:
            logging.info("Connected to MQTT broker successfully")
            # 订阅主题
            topic = self.config['mqtt']['subscribe_topic']
            client.subscribe(topic, qos=1)
            logging.info(f"Subscribed to topic: {topic}")
        else:
            logging.error(f"Failed to connect to MQTT broker, return code: {rc}")
    
    def _on_mqtt_disconnect(self, client, userdata, rc):
        """MQTT断开连接回调"""
        if rc != 0:
            logging.warning("Unexpected disconnection from MQTT broker")
    
    def _on_mqtt_message(self, client, userdata, msg):
        """MQTT接收消息回调"""
        topic = msg.topic
        payload = msg.payload
        
        self.logger.info(f"Received MQTT message from topic: {topic}, payload length: {len(payload)}")
        
        # 解析设备ID
        try:
            topic_parts = topic.split('/')
            device_id = topic_parts[self.device_id_position]
            
            # 查找或创建TCP连接
            tcp_conn = self._get_or_create_tcp_connection(device_id)
            
            if tcp_conn:
                # 转发消息到TCP
                try:
                    tcp_conn.sendall(payload)
                    self.logger.info(f"Forwarded message to device {device_id} via TCP")
                except Exception as e:
                    self.logger.error(f"Failed to send data to TCP for device {device_id}: {e}")
                    # 关闭连接并从映射表中移除
                    self._remove_tcp_connection(device_id)
        except Exception as e:
            self.logger.error(f"Error processing MQTT message: {e}")
    
    def _get_or_create_tcp_connection(self, device_id):
        """获取或创建TCP连接"""
        with self.mapping_lock:
            if device_id in self.tcp_connections:
                # 检查连接是否有效
                conn, last_active = self.tcp_connections[device_id]
                try:
                    # 测试连接是否仍然有效
                    conn.settimeout(0.1)
                    conn.sendall(b"")
                    # 更新最后活动时间
                    self.tcp_connections[device_id] = (conn, time.time())
                    return conn
                except:
                    # 连接无效，关闭并移除
                    self._remove_tcp_connection(device_id)
            
            # 创建新的TCP连接
            try:
                conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                conn.settimeout(self.tcp_connect_timeout)
                conn.connect((self.tcp_remote_host, self.tcp_remote_port))
                conn.settimeout(None)  # 设置为阻塞模式
                
                # 存储连接和最后活动时间
                self.tcp_connections[device_id] = (conn, time.time())
                
                # 创建反向映射（用于TCP数据到MQTT的转发）
                conn_id = id(conn)
                self.topic_to_device[conn_id] = device_id
                
                logging.info(f"Created TCP connection for device {device_id} to {self.tcp_remote_host}:{self.tcp_remote_port}")
                
                # 启动接收线程
                recv_thread = threading.Thread(target=self._tcp_recv_handler, args=(conn, device_id))
                recv_thread.daemon = True
                recv_thread.start()
                
                return conn
            except Exception as e:
                logging.error(f"Failed to create TCP connection for device {device_id}: {e}")
                return None
    
    def _tcp_recv_handler(self, conn, device_id):
        """TCP接收线程处理函数"""
        conn_id = id(conn)
        
        try:
            while self.running:
                # 接收TCP数据
                data = conn.recv(self.tcp_recv_buffer_size)
                
                if not data:
                    # 连接关闭
                    break
                
                # 更新最后活动时间
                with self.mapping_lock:
                    if device_id in self.tcp_connections:
                        self.tcp_connections[device_id] = (conn, time.time())
                
                # 检查是否是心跳包，如果是则不转发
                if self.heartbeat_enabled and data == self.heartbeat_packet:
                    self.logger.debug(f"Received heartbeat from device {device_id}, ignoring")
                    continue
                
                # 构建MQTT主题
                mqtt_topic = f"{self.config['mqtt']['publish_topic_prefix']}{device_id}"
                
                # 转发到MQTT
                if self.mqtt_client.is_connected():
                    self.mqtt_client.publish(mqtt_topic, data, qos=1)
                    self.logger.info(f"Forwarded TCP data from device {device_id} to MQTT topic: {mqtt_topic}")
                else:
                    self.logger.warning(f"MQTT client not connected, cannot forward data from device {device_id}")
        
        except Exception as e:
            self.logger.error(f"Error in TCP receive handler for device {device_id}: {e}")
        finally:
            # 移除连接
            with self.mapping_lock:
                self._remove_tcp_connection(device_id)
            
            logging.info(f"TCP connection closed for device {device_id}")
    
    def _remove_tcp_connection(self, device_id):
        """移除TCP连接"""
        if device_id in self.tcp_connections:
            conn, _ = self.tcp_connections[device_id]
            conn_id = id(conn)
            
            # 关闭连接
            try:
                conn.close()
            except:
                pass
            
            # 移除映射
            del self.tcp_connections[device_id]
            if conn_id in self.topic_to_device:
                del self.topic_to_device[conn_id]
                
    def _monitor_connections(self):
        """监控连接状态，用于GUI显示"""
        while True:
            time.sleep(10)
            
            with self.mapping_lock:
                active_connections = len(self.tcp_connections)
            
            mqtt_connected = False
            if self.mqtt_client:
                mqtt_connected = self.mqtt_client.is_connected()
            
            self.logger.debug(f"Connection status - TCP: {active_connections}, MQTT: {mqtt_connected}")
                
    def get_connection_stats(self):
        """获取连接统计信息，供GUI使用"""
        with self.mapping_lock:
            connections = list(self.tcp_connections.keys())
        
        mqtt_connected = False
        if self.mqtt_client:
            mqtt_connected = self.mqtt_client.is_connected()
        
        return {
            "tcp_connections": len(connections),
            "mqtt_connected": mqtt_connected,
            "devices": connections
        }
    
    def _heartbeat_thread(self):
        """心跳线程"""
        while self.running:
            time.sleep(self.heartbeat_interval)
            
            if not self.heartbeat_enabled:
                continue
            
            current_time = time.time()
            with self.mapping_lock:
                # 发送心跳并清理超时连接
                devices_to_remove = []
                
                for device_id, (conn, last_active) in self.tcp_connections.items():
                    # 检查连接是否超时
                    if current_time - last_active > self.config['mapping']['connection_timeout']:
                        devices_to_remove.append(device_id)
                        continue
                    
                    # 发送心跳包
                    try:
                        conn.sendall(self.heartbeat_packet)
                        # 更新最后活动时间
                        self.tcp_connections[device_id] = (conn, current_time)
                        logging.debug(f"Sent heartbeat to device {device_id}")
                    except Exception as e:
                        logging.error(f"Failed to send heartbeat to device {device_id}: {e}")
                        devices_to_remove.append(device_id)
                
                # 移除超时或失效的连接
                for device_id in devices_to_remove:
                    self._remove_tcp_connection(device_id)
                    logging.info(f"Removed inactive TCP connection for device {device_id}")
    
    def connect_mqtt(self):
        """连接MQTT broker"""
        try:
            mqtt_config = self.config['mqtt']
            self.mqtt_client.connect(
                mqtt_config['broker'],
                mqtt_config['port'],
                mqtt_config['keepalive']
            )
            return True
        except Exception as e:
            logging.error(f"Failed to connect to MQTT broker: {e}")
            return False
    
    def start(self):
        """启动网关"""
        self.running = True
        self.logger.info("TCP2MQTT Gateway starting...")
        
        # 启动心跳线程
        if self.heartbeat_enabled:
            self.heartbeat_thread = threading.Thread(target=self._heartbeat_thread)
            self.heartbeat_thread.daemon = True
            self.heartbeat_thread.start()
        
        # 连接MQTT
        if not self.connect_mqtt():
            # 如果连接失败，尝试重连
            reconnect_interval = 5
            while self.running:
                self.logger.info(f"Retrying to connect to MQTT broker in {reconnect_interval} seconds...")
                time.sleep(reconnect_interval)
                if self.connect_mqtt():
                    break
        
        # 启动MQTT循环
        try:
            self.mqtt_client.loop_forever()
        except KeyboardInterrupt:
            self.logger.info("Received keyboard interrupt, stopping gateway...")
        finally:
            self.stop()
    
    def stop(self):
        """停止网关"""
        self.running = False
        
        # 等待心跳线程结束
        if self.heartbeat_thread:
            self.heartbeat_thread.join(timeout=2.0)
        
        # 关闭所有TCP连接
        with self.mapping_lock:
            for device_id in list(self.tcp_connections.keys()):
                self._remove_tcp_connection(device_id)
        
        # 断开MQTT连接
        if self.mqtt_client:
            self.mqtt_client.disconnect()
            self.mqtt_client.loop_stop()
        
        logging.info("TCP2MQTT Gateway stopped")

def main():
    # 获取配置文件路径
    config_file = 'config.yaml'
    if len(sys.argv) > 1:
        config_file = sys.argv[1]
    
    # 创建并启动网关
    gateway = TCP2MQTTGateway(config_file)
    gateway.start()

if __name__ == '__main__':
    main()